#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define IN_FILENAME "in.ppm"
#define OUT_FILENAME "out.ppm"
#define N 10 // Number of threads
#define MASK_SIZE 5 // Must be odd

// Apply mask on a section of the image 
__global__ void apply_mask(unsigned int imgWidth, unsigned int imgHeight, unsigned char *imgIn, unsigned char *imgOut) {
    int tid = threadIdx.x;
  	unsigned int startRow, endRow;	
  	startRow = (tid*imgHeight)/N;
  	endRow =  ((tid+1)*imgHeight)/N;

  	int i, j, k, l;

  	for(i = startRow; i < endRow; i++){
  		for(j = 0; j < imgWidth; j++){
  			int totalR = 0;
  			int totalG = 0;
  			int totalB = 0;
  			for(k = -MASK_SIZE/2; k <= MASK_SIZE/2; k++){
  				for(l = -MASK_SIZE/2; l <= MASK_SIZE/2; l++){
  					if( (i + k >= 0) && (i + k < imgHeight) && (j + l >= 0) && (j + l < imgWidth) ){
  						totalR += imgIn[(i + k)*imgWidth*3 + (j + l) * 3];
  						totalG += imgIn[(i + k)*imgWidth*3 + (j + l) * 3 + 1];
  						totalB += imgIn[(i + k)*imgWidth*3 + (j + l) * 3 + 2];
  					}
  				}
  			}
  			imgOut[i*imgWidth*3 + j * 3] = totalR/(MASK_SIZE*MASK_SIZE);
  			imgOut[i*imgWidth*3 + j * 3 + 1] = totalG/(MASK_SIZE*MASK_SIZE);
  			imgOut[i*imgWidth*3 + j * 3 + 2] = totalB/(MASK_SIZE*MASK_SIZE);
  		}
  	}
}

int main(void)
{	
	unsigned char *imgIn, *imgOut;

	char magicNumber[3];

	unsigned int imgWidth, imgHeight, imgDepth;

	// Device Pointers
	unsigned char *dev_imgIn, *dev_imgOut;

	FILE *inFile;
	FILE *outFile;

	// Open file and read info
	inFile = fopen( IN_FILENAME , "r" );
	fscanf( inFile , "%s %u %u %u", magicNumber, &imgWidth, &imgHeight, &imgDepth);

	// Check image magic number
	if(strcmp("P3", magicNumber) != 0) {
		printf("Image magic number must be P3\n");
		return 0;
	}

	// allocate memory
	imgIn = (unsigned char*)malloc(imgWidth*imgHeight*3*sizeof(unsigned char));
	imgOut = (unsigned char*)malloc(imgWidth*imgHeight*3*sizeof(unsigned char));

	// Read image
	int i, j;
	for(i = 0; i < imgHeight; i++){
		for(j = 0; j < imgWidth; j++){
			fscanf(inFile, "%hhu", &(imgIn[imgWidth * i * 3 + j * 3]));
			fscanf(inFile, "%hhu", &(imgIn[imgWidth * i * 3 + j * 3 + 1]));
			fscanf(inFile, "%hhu", &(imgIn[imgWidth * i * 3 + j * 3 + 2]));
		}
	}

	// Close file
	fclose(inFile);

	// Allocate the memory on the GPU
    cudaMalloc( (void**)&dev_imgIn, imgWidth * imgHeight * 3 * sizeof(unsigned char));
    cudaMalloc( (void**)&dev_imgOut, imgWidth * imgHeight * 3 * sizeof(unsigned char));

    // Copy data to GPU
    cudaMemcpy( dev_imgIn, imgIn, imgWidth * imgHeight * 3 * sizeof(unsigned char), cudaMemcpyHostToDevice);

    // Call Kernel
    apply_mask<<<1, N>>>(imgWidth, imgHeight, dev_imgIn, dev_imgOut);

    // Copy output image from GPU
    cudaMemcpy(imgOut, dev_imgOut, imgWidth * imgHeight * 3 * sizeof(unsigned char), cudaMemcpyDeviceToHost);

    // Save output image to file
    outFile = fopen( OUT_FILENAME , "w" );
	fprintf( outFile , "P3\n%u %u\n %u\n", imgWidth, imgHeight, imgDepth);
	for(i = 0; i < imgHeight; i++){
		for(j = 0; j < imgWidth; j++){
			fprintf( outFile, "%hhu ", imgOut[imgWidth * i * 3 + j * 3]);
			fprintf( outFile, "%hhu ", imgOut[imgWidth * i * 3 + j * 3 + 1]);
			fprintf( outFile, "%hhu ", imgOut[imgWidth * i * 3 + j * 3 + 2]);
		}
		fprintf( outFile, "\n");
	}

    // Close file
    fclose(outFile);

    // Free allocated memory on GPU
    cudaFree(dev_imgIn);
    cudaFree(dev_imgOut);
    free(imgIn);
    free(imgOut);

	return 0;
}